package com.filter.control.common;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.*;

public class CustomServletRequestWrapper extends HttpServletRequestWrapper {

    private String modifiedUrl;
    
    private String method;

    public CustomServletRequestWrapper(HttpServletRequest request, String modifiedUrl) {
        super(request);
        this.modifiedUrl = modifiedUrl;
        this.method = request.getMethod();
        this.headers = new HashMap<>();
    }

    public CustomServletRequestWrapper(HttpServletRequest request, String modifiedUrl, String method) {
        super(request);
        this.modifiedUrl = modifiedUrl;
        this.method = method;
        this.headers = new HashMap<>();
    }

    @Override
    public String getRequestURI() {
        return modifiedUrl;
    }

    @Override
    public StringBuffer getRequestURL() {
        StringBuffer url = new StringBuffer(modifiedUrl);
        String queryString = getQueryString();
        if (queryString != null) {
            url.append("?").append(queryString);
        }
        return url;
    }

    @Override
    public String getMethod() {
        return method;
    }

    private Map<String, List<String>> headers;

    public void addHeader(String name, String value) {
        List<String> values = headers.getOrDefault(name, new ArrayList<>());
        values.add(value);
        headers.put(name, values);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        Set<String> headerNames = new HashSet<>(Collections.list(super.getHeaderNames()));
        headerNames.addAll(headers.keySet());
        return Collections.enumeration(headerNames);
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        List<String> headerValues = new ArrayList<>(Collections.list(super.getHeaders(name)));
        List<String> customHeaderValues = headers.getOrDefault(name, new ArrayList<>());
        headerValues.addAll(customHeaderValues);
        return Collections.enumeration(headerValues);
    }

    @Override
    public String getHeader(String name) {
        String headerValue = super.getHeader(name);
        if (headerValue == null) {
            List<String> customHeaderValues = headers.get(name);
            if (customHeaderValues != null && !customHeaderValues.isEmpty()) {
                headerValue = customHeaderValues.get(0);
            }
        }
        return headerValue;
    }
}